package org.database.config;

import com.mybatisflex.core.FlexGlobalConfig;
import com.mybatisflex.core.datasource.DataSourceBuilder;
import com.mybatisflex.core.datasource.DataSourceManager;
import com.mybatisflex.core.datasource.FlexDataSource;
import com.mybatisflex.core.dialect.DbTypeUtil;
import lombok.extern.slf4j.Slf4j;
import org.database.attribute.CustomDataSourceAttribute;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * DynamicDataSourceAutoConfiguration
 *
 * @author grl
 * @date 2024/2/2
 */
@Slf4j
public class CustomDataSourceConfiguration {

    private final String shardingDataSourceName;

    private static final String defaultDataSourceName = "MyBatis-Flex";

    private final Map<String, Map<String, String>> dataSourceProperties;

    public CustomDataSourceConfiguration(CustomDataSourceAttribute dataSourceProperties, String shardingDataSourceName) {
        this.dataSourceProperties = dataSourceProperties.getDatasource();
        this.shardingDataSourceName = shardingDataSourceName;
    }
    @SuppressWarnings("unchecked")
    public void createDataSource() {
        FlexDataSource flexDataSource = FlexGlobalConfig.getDefaultConfig().getDataSource();
        Map<String, DataSource> defaultDateSource = new HashMap<>(4);
        if (this.dataSourceProperties != null && !this.dataSourceProperties.isEmpty()) {
            for (Map.Entry<String, Map<String, String>> stringMapEntry : this.dataSourceProperties.entrySet()) {
                Map.Entry<String, Map<String, String>> entry = (Map.Entry) stringMapEntry;
                DataSource dataSource = (new DataSourceBuilder((Map) entry.getValue())).build();
                DataSourceManager.decryptDataSource((DataSource) dataSource);
                if (defaultDateSource.isEmpty()) {
                    defaultDateSource.put(entry.getKey(), (DataSource) dataSource);
                }
                if (flexDataSource == null) {
                    flexDataSource = new FlexDataSource((String) entry.getKey(), (DataSource) dataSource, false);
                } else {
                    flexDataSource.addDataSource((String) entry.getKey(), (DataSource) dataSource, false);
                }
            }
        }
        if (!defaultDateSource.isEmpty()) {
            Map<String, DataSource> dataSourceMap = flexDataSource.getDataSourceMap();
            DataSource dataSource = dataSourceMap.get(defaultDataSourceName);
            if (dataSource != null) {
                dataSourceMap.remove(defaultDataSourceName);
                dataSourceMap.put(shardingDataSourceName, dataSource);
            }
            try {
                String defaultDataSourceKey = defaultDateSource.keySet().iterator().next();
                DataSource defaultDataSource = defaultDateSource.get(defaultDataSourceKey);
                Class<? extends FlexDataSource> aClass = flexDataSource.getClass();
                Field dataSourceKeyField = aClass.getDeclaredField("defaultDataSourceKey");
                Field defaultDataSourceField = aClass.getDeclaredField("defaultDataSource");
                Field defaultDbType = aClass.getDeclaredField("defaultDbType");
                dataSourceKeyField.setAccessible(true);
                defaultDataSourceField.setAccessible(true);
                defaultDbType.setAccessible(true);
                dataSourceKeyField.set(flexDataSource, defaultDataSourceKey);
                defaultDataSourceField.set(flexDataSource, defaultDataSource);
                defaultDbType.set(flexDataSource, DbTypeUtil.getDbType(defaultDataSource));
            } catch (NoSuchFieldException | IllegalAccessException e) {
                log.error("修改默认数据源失败",e);
            }
        }
    }
}
